package cc.git.liuyan.customeraiagent.core.embeddingmodel;

import com.zaxxer.hikari.HikariConfig;
import com.zaxxer.hikari.HikariDataSource;
import lombok.extern.slf4j.Slf4j;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.text.PDFTextStripper;
import org.apache.poi.hwpf.HWPFDocument;
import org.apache.poi.hwpf.usermodel.Range;
import org.apache.poi.xwpf.usermodel.XWPFDocument;
import org.apache.poi.xwpf.usermodel.XWPFParagraph;

import java.io.*;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.*;

@Slf4j
public class SimpleEmbeddingModelSplit implements EmbeddingModelSplit {

    @Override
    public List<EmbeddingModelInputData> split(EmbeddingModelInputData inputData) {
        if (Objects.equals(inputData.getType(), "string")) {
            return SplitUtil.splitStringInputData(inputData);
        } else if (Objects.equals(inputData.getType(), "pdf")) {
            return this.splitPdfFile(inputData.getInputData());
        } else if (Objects.equals(inputData.getType(), "doc")) {
            return this.splitDocFile(inputData.getInputData());
        } else if (Objects.equals(inputData.getType(), "docx")) {
            return this.splitDocxFile(inputData.getInputData());
        } else if (Objects.equals(inputData.getType(), "txt")) {
            return this.splitTxtFile(inputData.getInputData());
        } else if (Objects.equals(inputData.getType(), "mysql")) {
            return this.splitMysql(((EmbeddingModelMysqlInputData) inputData).getOtherData());
        } else if (Objects.equals(inputData.getType(), "remark")) {
            //remark不分片
            return Collections.singletonList(new EmbeddingModelInputData("string", inputData.getInputData()));
        } else {
            throw new IllegalArgumentException("Unsupported type: " + inputData.getType());
        }
    }

    @Override
    public List<EmbeddingModelInputData> splitPdfFile(String pdfurl) {
        InputStream inputStream = UrlUtil.getUrlInputStream(pdfurl);
        List<EmbeddingModelInputData> inputDataList = splitPdfFile(inputStream);
        UrlUtil.closeUrlInputStream(inputStream);
        return inputDataList;
    }

    @Override
    public List<EmbeddingModelInputData> splitPdfFile(InputStream pdfInputStream) {
        try (PDDocument document = PDDocument.load(pdfInputStream)) {
            PDFTextStripper pdfStripper = new PDFTextStripper();
            String text = pdfStripper.getText(document);
            EmbeddingModelInputData inputData = new EmbeddingModelInputData(text);
            return SplitUtil.splitStringInputData(inputData);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public List<EmbeddingModelInputData> splitDocFile(String wordurl) {
        InputStream inputStream = UrlUtil.getUrlInputStream(wordurl);
        List<EmbeddingModelInputData> inputDataList = splitDocFile(inputStream);
        UrlUtil.closeUrlInputStream(inputStream);
        return inputDataList;
    }

    @Override
    public List<EmbeddingModelInputData> splitDocFile(InputStream wordInputStream) {
        try {
            HWPFDocument document = new HWPFDocument(wordInputStream);
            Range range = document.getRange();
            String text = range.text();
            EmbeddingModelInputData inputData = new EmbeddingModelInputData(text);
            return SplitUtil.splitStringInputData(inputData);
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
    }

    @Override
    public List<EmbeddingModelInputData> splitDocxFile(String wordurl) {
        InputStream inputStream = UrlUtil.getUrlInputStream(wordurl);
        List<EmbeddingModelInputData> inputDataList = splitDocxFile(inputStream);
        UrlUtil.closeUrlInputStream(inputStream);
        return inputDataList;
    }

    @Override
    public List<EmbeddingModelInputData> splitDocxFile(InputStream wordInputStream) {
        XWPFDocument document = null;
        try {
            StringBuilder contentBuilder = new StringBuilder();
            document = new XWPFDocument(wordInputStream);
            List<XWPFParagraph> paragraphs = document.getParagraphs();
            for (XWPFParagraph paragraph : paragraphs) {
                String text = paragraph.getText();
                contentBuilder.append(text);
            }
            EmbeddingModelInputData inputData = new EmbeddingModelInputData(contentBuilder.toString());
            return SplitUtil.splitStringInputData(inputData);
        } catch (IOException e) {
            throw new RuntimeException(e);
        } finally {
            try {
                if (document != null) {
                    document.close();
                }
            } catch (IOException e) {
                throw new RuntimeException(e);
            }
        }
    }

    @Override
    public List<EmbeddingModelInputData> splitTxtFile(String txturl) {
        InputStream inputStream = UrlUtil.getUrlInputStream(txturl);
        List<EmbeddingModelInputData> inputDataList = splitTxtFile(inputStream);
        UrlUtil.closeUrlInputStream(inputStream);
        return inputDataList;
    }

    @Override
    public List<EmbeddingModelInputData> splitTxtFile(InputStream txtInputStream) {
        List<EmbeddingModelInputData> inputDataList = new ArrayList<>();
        InputStreamReader inputStreamReader = new InputStreamReader(txtInputStream);
        BufferedReader bufferedReader = new BufferedReader(inputStreamReader);
        StringJoiner content = new StringJoiner("\n");
        String line = null;
        try {
            while ((line = bufferedReader.readLine()) != null) {
                if (line.trim().equals("")) {
                    EmbeddingModelInputData inputData = new EmbeddingModelInputData(content.toString());
                    inputDataList.addAll(SplitUtil.splitStringInputData(inputData));
                    content = new StringJoiner("\n");
                } else {
                    content.add(line);
                }
            }
            if (content.length() > 0) {
                EmbeddingModelInputData inputData = new EmbeddingModelInputData(content.toString());
                inputDataList.addAll(SplitUtil.splitStringInputData(inputData));
            }

            bufferedReader.close();
            inputStreamReader.close();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }

        return inputDataList;
    }

    @Override
    public List<EmbeddingModelInputData> splitMysql(MysqlInfo mysqlInfo) {
        List<EmbeddingModelInputData> segmentInputDataList = new ArrayList<>();
        HikariConfig config = new HikariConfig();
        config.setJdbcUrl(mysqlInfo.getUrl());
        config.setUsername(mysqlInfo.getUsername());
        config.setPassword(mysqlInfo.getPassword());
        String targetSchema = mysqlInfo.getTargetSchema();
        String targetTable = mysqlInfo.getTargetTable();
        try (HikariDataSource hikariDataSource = new HikariDataSource(config)) {
            try (Connection connection = hikariDataSource.getConnection()) {
                String sql = "select TABLE_SCHEMA, TABLE_NAME from information_schema.TABLES where TABLE_SCHEMA not in ('mysql', 'sys', 'information_schema')";
                if (targetSchema != null && !targetSchema.isEmpty() && targetTable != null && !targetTable.isEmpty() && !targetSchema.equals("*") && !targetTable.equals("*")) {
                    sql = String.format(sql.concat(" and TABLE_SCHEMA = '%s' and TABLE_NAME = '%s'"), targetSchema, targetTable);
                }
                try (PreparedStatement preparedStatement = connection.prepareStatement(sql)) {
                    ResultSet rowSet = preparedStatement.executeQuery();
                    while (rowSet.next()) {
                        //每张表单独转化
                        String tableSchema = rowSet.getString("TABLE_SCHEMA");
                        String tableName = rowSet.getString("TABLE_NAME");
                        String showddl = String.format("SHOW CREATE TABLE `%s`.`%s`", tableSchema, tableName);
                        try (PreparedStatement showddlPrepareStatement = connection.prepareStatement(showddl)) {
                            ResultSet showddlResultSet = showddlPrepareStatement.executeQuery();
                            if (showddlResultSet.next()) {
                                String tableddl = showddlResultSet.getString("Create Table");
                                String fullDDL = "CREATE TABLE `" + tableSchema + "`.`" + tableName + "` " + tableddl.substring(tableddl.indexOf("("));
                                segmentInputDataList.add(new EmbeddingModelInputData(fullDDL));
                            }
                        }
                    }
                }
            } catch (SQLException e) {
                throw new RuntimeException(e);
            }
        }
        return segmentInputDataList;
    }
}
